#include <mcp/server/fastmcp.h>
#include <mcp/server/tool_manager.h>
#include <mcp/server/resource_manager.h>
#include <mcp/server/prompt_manager.h>
#include <mcp/server/server_session.h>
#include <mcp/transport/transport.h>
#include <mcp/transport/stdio_transport.h>
#include <mcp/json/value.h>
#include <mcp/error.h>
#include <mcp/types.h>
#include <iostream>
#include <sstream>
#include <memory>
#include <algorithm>
#include <random>
#include <chrono>

namespace mcp {
namespace server {

namespace detail {
    // Context implementation for server sessions
    template<typename ServerSessionT, typename LifespanT>
    class ContextImpl {
    public:
        ContextImpl(std::weak_ptr<ServerSessionT> session,
                   std::shared_ptr<LifespanT> lifespan_context,
                   types::RequestId request_id)
            : session_(std::move(session))
            , lifespan_context_(std::move(lifespan_context))
            , request_id_(std::move(request_id)) {}

        // Logging implementation
        void debug(const std::string& message) {
            if (auto session = session_.lock()) {
                session->log(types::LoggingMessage{
                    .level = types::LogLevel::Debug,
                    .message = message
                });
            }
        }

        void info(const std::string& message) {
            if (auto session = session_.lock()) {
                session->log(types::LoggingMessage{
                    .level = types::LogLevel::Info,
                    .message = message
                });
            }
        }

        void warning(const std::string& message) {
            if (auto session = session_.lock()) {
                session->log(types::LoggingMessage{
                    .level = types::LogLevel::Warning,
                    .message = message
                });
            }
        }

        void error(const std::string& message) {
            if (auto session = session_.lock()) {
                session->log(types::LoggingMessage{
                    .level = types::LogLevel::Error,
                    .message = message
                });
            }
        }

        // Progress reporting implementation
        void report_progress(double progress,
                            std::optional<double> total,
                            std::optional<std::string> message) {
            if (auto session = session_.lock()) {
                session->send_progress(types::ProgressNotification{
                    .token = request_id_,
                    .progress = progress,
                    .total = total,
                    .message = std::move(message)
                });
            }
        }

        // Resource access implementation
        Result<types::ResourceContent> read_resource(const std::string& uri) {
            if (auto session = session_.lock()) {
                return session->read_resource(uri);
            }
            return Error::internal_error("Session expired");
        }

        // User interaction implementation
        template<typename SchemaT>
        Result<types::ElicitationResult> elicit(const std::string& message, const SchemaT& schema) {
            if (auto session = session_.lock()) {
                // This is a simplified implementation
                return Error::unsupported_operation("Elicitation not implemented");
            }
            return Error::internal_error("Session expired");
        }

        // LLM sampling implementation
        Result<types::SamplingResult> create_message(
            const std::vector<types::SamplingMessage>& messages,
            const types::SamplingOptions& options) {
            if (auto session = session_.lock()) {
                return session->create_message(messages, options);
            }
            return Error::internal_error("Session expired");
        }

        // Context access implementation
        ServerSessionT& session() {
            if (auto session = session_.lock()) {
                return *session;
            }
            throw std::runtime_error("Session expired");
        }

        LifespanT& lifespan_context() {
            if (!lifespan_context_) {
                throw std::runtime_error("No lifespan context available");
            }
            return *lifespan_context_;
        }

    private:
        std::weak_ptr<ServerSessionT> session_;
        std::shared_ptr<LifespanT> lifespan_context_;
        types::RequestId request_id_;
    };
}

// FastMCPServer implementation
struct FastMCPServer::Impl {
    ServerConfig config;
    std::unique_ptr<ToolManager> tool_manager;
    std::unique_ptr<ResourceManager> resource_manager;
    std::unique_ptr<PromptManager> prompt_manager;
    std::unique_ptr<LifespanManager<void>> lifespan_manager;

    // Server state
    std::vector<std::weak_ptr<ServerSession>> active_sessions;
    mutable std::mutex sessions_mutex;

    // UUID generation
    std::random_device rd;
    std::mt19937 gen;
    std::uniform_int_distribution<uint64_t> dis;

    Impl(const ServerConfig& config)
        : config(config)
        , tool_manager(std::make_unique<ToolManager>())
        , resource_manager(std::make_unique<ResourceManager>())
        , prompt_manager(std::make_unique<PromptManager>())
        , lifespan_manager(std::make_unique<LifespanManager<void>>())
        , rd()
        , gen(rd())
        , dis(0, UINT64_MAX) {}
};

FastMCPServer::FastMCPServer(const ServerConfig& config)
    : impl_(std::make_unique<Impl>(config)) {}

FastMCPServer::~FastMCPServer() = default;

// Server information
const ServerConfig& FastMCPServer::config() const {
    return impl_->config;
}

std::vector<types::Tool> FastMCPServer::list_tools() const {
    return impl_->tool_manager->list_tools();
}

std::vector<types::Resource> FastMCPServer::list_resources() const {
    return impl_->resource_manager->list_resources();
}

std::vector<types::ResourceTemplate> FastMCPServer::list_resource_templates() const {
    return impl_->resource_manager->list_resource_templates();
}

std::vector<types::Prompt> FastMCPServer::list_prompts() const {
    return impl_->prompt_manager->list_prompts();
}

// Session management
std::shared_ptr<ServerSession> FastMCPServer::create_session() {
    auto session_id = generate_session_id();
    auto transport = std::make_unique<transport::StdioTransport>();

    auto session = std::make_shared<ServerSession>(
        std::move(transport),
        impl_->tool_manager,
        impl_->resource_manager,
        impl_->prompt_manager
    );

    // Add to active sessions
    {
        std::lock_guard<std::mutex> lock(impl_->sessions_mutex);
        impl_->active_sessions.push_back(session);
    }

    // Set up session cleanup
    session->set_close_handler([this, session]() {
        std::lock_guard<std::mutex> lock(impl_->sessions_mutex);
        impl_->active_sessions.erase(
            std::remove_if(impl_->active_sessions.begin(), impl_->active_sessions.end(),
                [session](const std::weak_ptr<ServerSession>& weak) {
                    return weak.expired() || weak.lock() == session;
                }),
            impl_->active_sessions.end()
        );
    });

    return session;
}

// Tool registration - template implementation
template<typename Func>
void FastMCPServer::tool(const std::string& name, Func&& func, const ToolOptions& options) {
    // Extract function signature and create appropriate handler
    if constexpr (std::is_invocable_v<Func, const json::Value&, Context<>&>) {
        // Function takes JSON arguments and Context
        using ResultType = std::invoke_result_t<Func, const json::Value&, Context<>&>;

        if constexpr (std::is_same_v<ResultType, Result<json::Value>>) {
            // Function returns Result<json::Value>
            auto handler = [func = std::forward<Func>(func)](const json::Value& args, Context<>& context) -> Result<ToolResult> {
                auto result = func(args, context);
                if (!result) {
                    return result.error();
                }
                return ToolResult{
                    .content = {std::make_shared<types::TextContent>(result->serialize())},
                    .is_error = false,
                    .structured_content = *result
                };
            };

            ToolRegistration registration{
                .name = name,
                .description = options.description.value_or("Tool: " + name),
                .input_schema = options.input_schema.value_or(json::object_schema()),
                .output_schema = options.output_schema,
                .annotations = options.annotations,
                .handler = std::move(handler)
            };

            impl_->tool_manager->register_tool(std::move(registration));
        } else {
            // Function returns something else - try to convert to JSON
            auto handler = [func = std::forward<Func>(func)](const json::Value& args, Context<>& context) -> Result<ToolResult> {
                try {
                    auto result = func(args, context);
                    return ToolResult{
                        .content = {std::make_shared<types::TextContent>(json::Value(result).serialize())},
                        .is_error = false
                    };
                } catch (const std::exception& e) {
                    return Error::internal_error(std::string("Tool execution failed: ") + e.what());
                }
            };

            ToolRegistration registration{
                .name = name,
                .description = options.description.value_or("Tool: " + name),
                .input_schema = options.input_schema.value_or(json::object_schema()),
                .output_schema = options.output_schema,
                .annotations = options.annotations,
                .handler = std::move(handler)
            };

            impl_->tool_manager->register_tool(std::move(registration));
        }
    } else {
        // Function takes different arguments - try to adapt
        auto handler = [func = std::forward<Func>(func)](const json::Value& args, Context<>& context) -> Result<ToolResult> {
            try {
                // Try to call with empty arguments
                if constexpr (std::is_invocable_v<Func, Context<>&>) {
                    auto result = func(context);
                    return ToolResult{
                        .content = {std::make_shared<types::TextContent>(json::Value(result).serialize())},
                        .is_error = false
                    };
                } else if constexpr (std::is_invocable_v<Func>) {
                    auto result = func();
                    return ToolResult{
                        .content = {std::make_shared<types::TextContent>(json::Value(result).serialize())},
                        .is_error = false
                    };
                } else {
                    return Error::internal_error("Unsupported function signature for tool");
                }
            } catch (const std::exception& e) {
                return Error::internal_error(std::string("Tool execution failed: ") + e.what());
            }
        };

        ToolRegistration registration{
            .name = name,
            .description = options.description.value_or("Tool: " + name),
            .input_schema = options.input_schema.value_or(json::object_schema()),
            .output_schema = options.output_schema,
            .annotations = options.annotations,
            .handler = std::move(handler)
        };

        impl_->tool_manager->register_tool(std::move(registration));
    }
}

// Resource registration - template implementation
template<typename Func>
void FastMCPServer::resource(const std::string& uri_template, Func&& func, const ResourceOptions& options) {
    auto handler = [func = std::forward<Func>(func)](const std::string& uri, Context<>& context) -> Result<types::ResourceContent> {
        try {
            return func(uri, context);
        } catch (const std::exception& e) {
            return Error::internal_error(std::string("Resource handler failed: ") + e.what());
        }
    };

    ResourceRegistration registration{
        .uri_template = uri_template,
        .name = options.name.value_or("Resource: " + uri_template),
        .description = options.description,
        .mime_type = options.mime_type,
        .handler = std::move(handler)
    };

    impl_->resource_manager->register_resource(std::move(registration));
}

// Prompt registration - template implementation
template<typename Func>
void FastMCPServer::prompt(const std::string& name, Func&& func, const PromptOptions& options) {
    auto handler = [func = std::forward<Func>(func)](const json::Value& args, Context<>& context) -> Result<std::vector<types::PromptMessage>> {
        try {
            return func(args, context);
        } catch (const std::exception& e) {
            return Error::internal_error(std::string("Prompt handler failed: ") + e.what());
        }
    };

    // Default to no arguments - user can customize as needed
    std::vector<PromptArgumentDefinition> arguments;

    PromptRegistration registration{
        .name = name,
        .description = options.description,
        .arguments = std::move(arguments),
        .handler = std::move(handler)
    };

    impl_->prompt_manager->register_prompt(std::move(registration));
}

// Lifespan management
template<typename LifespanT, typename Func>
void FastMCPServer::lifespan(Func&& func) {
    static_assert(std::is_invocable_v<Func, LifespanT&>,
                 "Lifespan function must be invocable with LifespanT&");

    auto lifespan_mgr = std::make_unique<LifespanManager<LifespanT>>();
    lifespan_mgr->on_startup(std::forward<Func>(func));
    impl_->lifespan_manager = std::move(lifespan_mgr);
}

// Helper functions
std::string FastMCPServer::generate_session_id() const {
    std::stringstream ss;
    ss << std::hex << impl_->dis(impl_->gen);
    return ss.str();
}

void FastMCPServer::cleanup_expired_sessions() {
    std::lock_guard<std::mutex> lock(impl_->sessions_mutex);
    impl_->active_sessions.erase(
        std::remove_if(impl_->active_sessions.begin(), impl_->active_sessions.end(),
            [](const std::weak_ptr<ServerSession>& weak) { return weak.expired(); }),
        impl_->active_sessions.end()
    );
}

size_t FastMCPServer::get_active_session_count() const {
    std::lock_guard<std::mutex> lock(impl_->sessions_mutex);
    return std::count_if(impl_->active_sessions.begin(), impl_->active_sessions.end(),
        [](const std::weak_ptr<ServerSession>& weak) { return !weak.expired(); });
}

// Server capabilities
json::Value FastMCPServer::get_capabilities() const {
    return json::Value::object({
        {"tools", json::Value::object({
            {"list_changed", true},
            {"call", true}
        })},
        {"resources", json::Value::object({
            {"subscribe", true},
            {"list_changed", true},
            {"read", true}
        })},
        {"prompts", json::Value::object({
            {"list_changed", true},
            {"call", true}
        })},
        {"logging", json::Value::object({
            {"level", "debug"}
        })},
        {"progress", true}
    });
}

// Server info
json::Value FastMCPServer::get_server_info() const {
    return json::Value::object({
        {"name", impl_->config.name},
        {"version", impl_->config.version},
        {"protocolVersion", "2025-06-18"},
        {"capabilities", get_capabilities()},
        {"serverInfo", json::Value::object({
            {"name", impl_->config.name},
            {"version", impl_->config.version}
        })}
    });
}

// Explicit template instantiations
template void FastMCPServer::tool<std::function<Result<json::Value>(const json::Value&, Context<>&)>>(
    const std::string& name, std::function<Result<json::Value>(const json::Value&, Context<>&)>&& func, const ToolOptions& options);

template void FastMCPServer::resource<std::function<Result<types::ResourceContent>(const std::string&, Context<>&)>>(
    const std::string& uri_template, std::function<Result<types::ResourceContent>(const std::string&, Context<>&)>&& func, const ResourceOptions& options);

template void FastMCPServer::prompt<std::function<Result<std::vector<types::PromptMessage>>(const json::Value&, Context<>&)>>(
    const std::string& name, std::function<Result<std::vector<types::PromptMessage>>(const json::Value&, Context<>&)>&& func, const PromptOptions& options);

template void FastMCPServer::lifespan<void, std::function<void(void&)>>(
    std::function<void(void&)>&& func);

} // namespace server
} // namespace mcp